package at.ltasat.train.set;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

import at.ltasat.classifier.Clazz;
import at.ltasat.classifier.ClazzManager;
import at.ltasat.classifier.ObjectFeatures;
import at.ltasat.util.IO;

public final class TrainingSet<D> {

	private final int nDimensions;
	private final ClazzManager<D> clazzManager;
	private final Map<ObjectFeatures, int[]> data;

	public TrainingSet(int nDimensions, ClazzManager<D> clazzManager) {
		this.nDimensions = nDimensions;
		this.clazzManager = clazzManager;
		data = new HashMap<ObjectFeatures, int[]>();
	}

	public static <D> TrainingSet<D> read(final InputStream in, final ClazzManager<D> clazzManager) throws IOException {
		TrainingSet<D> set = new TrainingSet<D>(IO.readInteger(in), clazzManager);
		final int size = IO.readInteger(in);
		for (int count = 0; count < size; count++) {
			final double[] features = new double[set.getNumberOfDimensions()];
			for (int i = 0; i < features.length; i++) {
				features[i] = IO.readDouble(in);
			}
			final int[] counts = new int[clazzManager.size()];
			for (int i = 0; i < clazzManager.size(); i++) {
				counts[i] = IO.readInteger(in);
			}
			set.data.put(new ObjectFeatures(features), counts);
		}
		return set;
	}

	public void write(final OutputStream out) throws IOException {
		IO.writeInteger(out, getNumberOfDimensions());
		IO.writeInteger(out, data.size());
		for (ObjectFeatures point : getPoints()) {
			for (int i = 0; i < nDimensions; i++) {
				IO.writeDouble(out, point.getFeatureValue(i));
			}
			for (Clazz<D> clazz : clazzManager) {
				IO.writeInteger(out, getClassificationCount(point, clazz));
			}
		}
	}

	public int getNumberOfDimensions() {
		return nDimensions;
	}

	public Set<ObjectFeatures> getPoints() {
		return data.keySet();
	}

	public int[] getClassifications(ObjectFeatures point) {
		return data.get(point);
	}

	public int getClassificationCount(ObjectFeatures point, Clazz<D> clazz) {
		if (point.getNumberOfFeatures() != nDimensions) {
			throw new IllegalArgumentException("Number of dimensions don't match!");
		}
		int count;
		if (data.containsKey(point)) {
			return data.get(point)[clazz.getId()];
		} else {
			count = -1;
		}
		return count;
	}

	public void add(ObjectFeatures point, Clazz<D> clazz) {
		if (!data.containsKey(point)) {
			data.put(point, new int[clazzManager.size()]);
		}
		data.get(point)[clazz.getId()]++;
	}
}
